import math
import numpy as np
import pickle
from scipy import interpolate
import matplotlib as mpl
import matplotlib.pyplot as plt
from scipy.optimize import minimize
from scipy.optimize import brute
from matplotlib.ticker import (MultipleLocator, AutoMinorLocator)

plt.rcParams.update({'font.size': 20})
mpl.rcParams['axes.formatter.useoffset'] = False
spectral_dict=pickle.load(open("spectral_dict_stacked.pickle", "rb"))
spectral_dict_unstack=pickle.load(open("spectral_dict.pickle", "rb"))
#spectral_template=pickle.load(open("spectral_template_lime.pickle", "rb"))
#spectral_template_coms=pickle.load(open("spectral_template.pickle", "rb"))

fitline='HDO_225'

restfreq_hdo=225.89672e9
restfreq_ch3ocho=225.90074e9
restfreq_13ch3cho=225.89294840e9
linefreqs=[restfreq_hdo,restfreq_ch3ocho,restfreq_13ch3cho]
lines=['           HDO','    CH$_3$OCHO','$^{13}$CH$_3$CHO']
lineamp=[0.26,0.225,0.1]

freq_axis=spectral_dict[fitline]['freq_axis'].copy()
vel_axis=(freq_axis-restfreq_hdo)/restfreq_hdo*3.0e5
spectrum=spectral_dict[fitline]['spectrum'].copy()
e_spectrum=spectral_dict[fitline]['e_spectrum'].copy()


fig,ax=plt.subplots(1,1,figsize=(15, 10))
ax.plot(freq_axis,spectrum,drawstyle='steps',linewidth='2',color='black',label='Full Spectrum')

ax.xaxis.set_major_locator(MultipleLocator(0.005))
ax.xaxis.set_minor_locator(AutoMinorLocator())
ax.tick_params(which='both', width=2)
ax.tick_params(which='major', length=8)
ax.tick_params(which='minor', length=4)
ax.set_ylabel('Flux Density (Jy)')
ax.set_xlabel('Frequency (GHz)')
ax.set_xlim(np.min(spectral_dict_unstack[fitline]['freq_axis']/1.0e9),np.max(spectral_dict_unstack[fitline]['freq_axis']/1.0e9))
ax.set_ylim(np.min(spectrum)-0.05*np.min(spectrum),np.max(spectrum)+0.2*np.max(spectrum))
for i in range(len(lines)):
   ax.text(linefreqs[i]/1e9-(4.3+3.0)/3.0e5*linefreqs[i]/1e9,lineamp[i]+0.015,lines[i])
   ax.text(linefreqs[i]/1e9-4.3/3.0e5*linefreqs[i]/1e9,lineamp[i],'|')
ax.text(plt.xlim()[1]-plt.xlim()[1]*0.000006,plt.ylim()[1]-plt.ylim()[1]*0.05,'b)')
plt.savefig('HDO_225_stack_line_ids.png')
plt.savefig('HDO_225_stack_line_ids.pdf')

fitline='HDO_241'

restfreq_hdo=241.56155e9
restfreq_ch3cho=241.56322770e9
restfreq_ch32co=241.55988690e9
linefreqs=[restfreq_hdo,restfreq_ch3cho,restfreq_ch32co]
lines=['             HDO','             CH$_3$CHO','(CH$_3$)$_2$CO']
lineamp=[0.275,0.19,0.1]

freq_axis=spectral_dict[fitline]['freq_axis'].copy()
vel_axis=(freq_axis-restfreq_hdo)/restfreq_hdo*3.0e5
spectrum=spectral_dict[fitline]['spectrum'].copy()
e_spectrum=spectral_dict[fitline]['e_spectrum'].copy()


fig,ax=plt.subplots(1,1,figsize=(15, 10))
ax.plot(freq_axis,spectrum,drawstyle='steps',linewidth='2',color='black',label='Full Spectrum')

ax.xaxis.set_major_locator(MultipleLocator(0.005))
ax.xaxis.set_minor_locator(AutoMinorLocator())
ax.tick_params(which='both', width=2)
ax.tick_params(which='major', length=8)
ax.tick_params(which='minor', length=4)
ax.set_ylabel('Flux Density (Jy)')
ax.set_xlabel('Frequency (GHz)')
ax.set_xlim(np.min(spectral_dict_unstack[fitline]['freq_axis']/1.0e9),np.max(spectral_dict_unstack[fitline]['freq_axis']/1.0e9))
ax.set_ylim(np.min(spectrum)-0.05*np.min(spectrum),np.max(spectrum)+0.2*np.max(spectrum))
for i in range(len(lines)):
   ax.text(linefreqs[i]/1e9-(4.3+3.0)/3.0e5*linefreqs[i]/1e9,lineamp[i]+0.015,lines[i])
   ax.text(linefreqs[i]/1e9-4.3/3.0e5*linefreqs[i]/1e9,lineamp[i],'|')
ax.text(plt.xlim()[1]-plt.xlim()[1]*0.000006,plt.ylim()[1]-plt.ylim()[1]*0.05,'d)')
plt.savefig('HDO_241_stack_line_ids.png')
plt.savefig('HDO_241_stack_line_ids.pdf')



fitline='H218O_203'
restfreq_h218o=203.40752000e9
restfreq_hdo=restfreq_h218o
restfreq_ch3och3_1=203.41011260e9
restfreq_ch3och3_2=203.41143090e9
restfreq_ch3och3_3=203.40277790e9
linefreqs=[restfreq_h218o,restfreq_ch3och3_1,restfreq_ch3och3_2,restfreq_ch3och3_3]
lines=['H$_2$$^{18}$O','    CH$_3$OCH$_3$','','CH$_3$OCH$_3$']
lineamp=[0.06,0.15,0.15,0.09]

freq_axis=spectral_dict[fitline]['freq_axis'].copy()
vel_axis=(freq_axis-restfreq_hdo)/restfreq_hdo*3.0e5
spectrum=spectral_dict[fitline]['spectrum'].copy()
e_spectrum=spectral_dict[fitline]['e_spectrum'].copy()


fig,ax=plt.subplots(1,1,figsize=(15, 10))
ax.plot(freq_axis,spectrum,drawstyle='steps',linewidth='2',color='black',label='Full Spectrum')

ax.xaxis.set_major_locator(MultipleLocator(0.005))
ax.xaxis.set_minor_locator(AutoMinorLocator())
ax.tick_params(which='both', width=2)
ax.tick_params(which='major', length=8)
ax.tick_params(which='minor', length=4)
ax.set_ylabel('Flux Density (Jy)')
ax.set_xlabel('Frequency (GHz)')
ax.set_xlim(np.min(spectral_dict_unstack[fitline]['freq_axis']/1.0e9),np.max(spectral_dict_unstack[fitline]['freq_axis']/1.0e9))
ax.set_ylim(np.min(spectrum)-0.05*np.min(spectrum),np.max(spectrum)+0.2*np.max(spectrum))
for i in range(len(lines)):
   ax.text(linefreqs[i]/1e9-(4.3+1.5)/3.0e5*linefreqs[i]/1e9,lineamp[i]+0.01,lines[i])
   ax.text(linefreqs[i]/1e9-4.3/3.0e5*linefreqs[i]/1e9,lineamp[i],'|')
ax.text(plt.xlim()[1]-plt.xlim()[1]*0.000006,plt.ylim()[1]-plt.ylim()[1]*0.05,'f)')
plt.savefig('H218O_203_stack_line_ids.png')
plt.savefig('H218O_203_stack_line_ids.pdf')

